import matplotlib
matplotlib.use('Agg')
import copy
import datetime
import os
from utils.options import args_parser
from models.Update import *
# from models.resnet import *
from models.ResNet_Dery import *
from models.test import *
from models.aggregation import *
from models.branchnet import *
from models.cnn import *
from models.ResNet8 import *
# from models.ResNet20 import *
from models.MobileNetV2 import *
from models.MobileNetV3 import *
from models.Vgg16 import *
from models.ViT_T import *
from models.LSTM import *
from algorithm.Training_FedAvg import *

from algorithm.Training_FedDF2 import *

# from algorithm.Training_FedDery import *
# from algorithm.Training_FedDery2 import *
# from algorithm.Training_FedDery3 import *
from algorithm.Training_FedDery4 import *
from algorithm.Training_Base import *

from algorithm.Training_ScaleFL import ScaleFL
from algorithm.Training_HeteroFL import HeteroFL
from algorithm.Training_DepthFL import DepthFL
from utils.get_dataset import get_dataset
from utils.save_result import save_result
from utils.set_seed import set_random_seed
import wandb


if __name__ == '__main__':
    args = args_parser()
    args.device = torch.device('cuda:{}'.format(args.device  ) if torch.cuda.is_available() and args.device != -1 else 'cpu')
    print(torch.cuda.is_available())
    set_random_seed(args.seed)

    dataset_train,dataset_test,dict_users ,dict_global= get_dataset(args)
    net_list = []
    # os.environ["WANDB_API_KEY"] = "personal key"
    print(args.algorithm + '_' + args.dataset+'_'+args.model+'_'+str(args.iid)+'_'+str(args.noniid_case)+'_'+str(args.data_beta))
    # wandb.init(project="FedDery2",name='limit_{}_{}_{}_{}_{}_d{}_{}_{}_{}_lr_{}_id{}_id{}_block3_client_442_limit_ratio{}-{}-{}_layers_2_{}'.format(args.dataset,args.num_users,args.frac,args.iid, args.noniid_case,args.data_beta,args.algorithm, args.model,
    #                                                              args.epochs, args.lr, args.num_id1, args.num_id2,args.client1_frac,args.client2_frac,args.client3_frac,datetime.datetime.now().strftime("%Y_%m_%d_%H_%M_%S")),
    #                                                              tags=["{}".format(args.dataset),"{}".format(args.algorithm),"iid{}".format(args.iid),"{}".format(args.model)])
    if args.algorithm == 'FedAvg':
        
        # net_glob = MobileNetV3(num_classes = args.num_classes) 
        # net_glob = ViT_T12_cifar10(num_classes = args.num_classes)
        # net_glob = CNNCifar100(args=args)
        # net_glob = ResNet18_cifar10(num_classes = args.num_classes,args=args)
        # net_glob = mobilenetv2(args)
        # net_glob = VGG16(args)
        if args.model == 'cnn2':
            net_glob = CNNCifar100(args=args)
            temp = 0
        elif args.model == 'mobilenetv2':
            net_glob = mobilenetv2(args)
            temp = 1
        elif args.model == 'resnet18':
            net_glob = ResNet18_cifar10(num_classes = args.num_classes,args=args)
            temp = 2
        else:
            exit("model error")
        # net_glob = CharLSTM()
        # temp = 0
        print(net_glob)
        print(temp)
        net_glob.to(args.device)
        FedAvg(args,net_glob,dataset_train,dataset_test,dict_users,temp)


    elif args.algorithm == "FedDF":
        net_zoo = []
        if args.model == 'cnn_mobilenetv2_resnet18':
            net1 = CNNCifar100(args=args)
            net2 = mobilenetv2(args)
            net3 = ResNet18_cifar10(num_classes = args.num_classes,args=args)
        elif args.model == 'mobilenetv2_resnet18_vgg16':
            net1 = mobilenetv2(args)
            net2 = ResNet18_cifar10(num_classes = args.num_classes,args=args)
            net3 = VGG16(args)
        else:
            exit("model error")
        net_zoo.append(net1.to(args.device))
        net_zoo.append(net2.to(args.device))
        # net3 = VGG16(args)
        # net3 = ResNet50_cifar10(num_classes = args.num_classes)
        net_zoo.append(net3.to(args.device))
        FedDF(args,net_zoo,dataset_train,dataset_test,dict_users,dict_global)
 

    elif args.algorithm == "FedDery":
        net_zoo = []
        # net1 = ResNet18_cifar10(num_classes = args.num_classes)
        # net2 = ResNet50_cifar10(num_classes = args.num_classes)
        # net_zoo.append(net1.to(args.device))
        # net_zoo.append(net2.to(args.device))
        # net3 = MobileNetV2(num_classes = args.num_classes)
        # net_zoo.append(net3.to(args.device))
        if args.model == 'cnn_mobilenetv2_resnet18':
            net1 = CNNCifar100(args=args)
            net2 = mobilenetv2(args)
            net3 = ResNet18_cifar10(num_classes = args.num_classes,args=args)
        elif args.model == 'mobilenetv2_resnet18_vgg16':
            net1 = mobilenetv2(args)
            net2 = ResNet18_cifar10(num_classes = args.num_classes,args=args)
            net3 = VGG16(args)
        else:
            exit("model error")
        net_zoo.append(net1.to(args.device))
        net_zoo.append(net2.to(args.device))
        # net3 = VGG16(args)
        # net3 = ResNet50_cifar10(num_classes = args.num_classes)
        net_zoo.append(net3.to(args.device))
        FedDery(args,net_zoo,dataset_train,dataset_test,dict_users,dict_global)
    
    elif args.algorithm == "FedBase":
        net_zoo = []
        # net1 = ResNet18_cifar10(num_classes = args.num_classes)
        # net2 = ResNet50_cifar10(num_classes = args.num_classes)
        # net_zoo.append(net1.to(args.device))
        # net_zoo.append(net2.to(args.device))
        # # net1 = ResNet8(num_classes = args.num_classes)
        # # net2 = ResNet18_cifar10(num_classes = args.num_classes)
        # # net_zoo.append(net1.to(args.device))
        # # net_zoo.append(net2.to(args.device))
        # net3 = MobileNetV3(num_classes = args.num_classes)
        # net_zoo.append(net3.to(args.device))
        if args.model == 'cnn_mobilenetv2_resnet18':
            net1 = CNNCifar100(args=args)
            net2 = mobilenetv2(args)
            net3 = ResNet18_cifar10(num_classes = args.num_classes,args=args)
        elif args.model == 'mobilenetv2_resnet18_vgg16':
            net1 = mobilenetv2(args)
            net2 = ResNet18_cifar10(num_classes = args.num_classes,args=args)
            net3 = VGG16(args)
        else:
            exit("model error")
        net_zoo.append(net1.to(args.device))
        net_zoo.append(net2.to(args.device))
        # net3 = VGG16(args)
        # net3 = ResNet50_cifar10(num_classes = args.num_classes)
        net_zoo.append(net3.to(args.device))
        FedBase(args,net_zoo,dataset_train,dataset_test,dict_users)

    elif args.algorithm == "HeteroFL":

        HeteroFL(args,dataset_train,dataset_test,dict_users)

    elif args.algorithm == "DepthFL":

        DepthFL(args,dataset_train,dataset_test,dict_users)

    elif args.algorithm == "FedTest":
        from algorithm.Training_FedTest import FedTest

        net_zoo = []
        # net1 = ResNet18_cifar10(num_classes = args.num_classes)
        # net2 = ResNet50_cifar10(num_classes = args.num_classes)
        # net_zoo.append(net1.to(args.device))
        # net_zoo.append(net2.to(args.device))
        # net3 = MobileNetV2(num_classes = args.num_classes)
        # net_zoo.append(net3.to(args.device))
        if args.model == 'cnn_mobilenetv2_resnet18':
            net1 = CNNCifar100(args=args)
            net2 = mobilenetv2(args)
            net3 = ResNet18_cifar10(num_classes=args.num_classes, args=args)
        elif args.model == 'mobilenetv2_resnet18_vgg16':
            net1 = mobilenetv2(args)
            net2 = ResNet18_cifar10(num_classes=args.num_classes, args=args)
            net3 = VGG16(args)
        else:
            exit("model error")
        net_zoo.append(net1.to(args.device))
        net_zoo.append(net2.to(args.device))
        # net3 = VGG16(args)
        # net3 = ResNet50_cifar10(num_classes = args.num_classes)
        net_zoo.append(net3.to(args.device))
        FedTest(args, net_zoo, dataset_train, dataset_test, dict_users, dict_global)

    elif args.algorithm == "ScaleFL":
        ScaleFL(args, dataset_train, dataset_test, dict_users)


